{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predicting logP for JAK2 with RNNs and SMILES strings "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates how to build predictive recurrent neural network for SMILES strings. We will build classification model for logP with OpenChem Toolkit (https://github.com/Mariewelt/OpenChem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fatal: destination path 'OpenChem' already exists and is not an empty directory.\r\n"
     ]
    }
   ],
   "source": [
    "# Cloning OpenChem. Comment this line if you already cloned the repository\n",
    "! git clone --branch develop https://github.com/Mariewelt/OpenChem.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: CUDA_VISIBLE_DEVICES=2\n"
     ]
    }
   ],
   "source": [
    "%env CUDA_VISIBLE_DEVICES=2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/home/mariewelt/Work/tmp/OpenChem'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sys.path.pop(sys.path.index('/home/mariewelt/Work/tmp/OpenChem'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.append('/home/mariewelt/Work/tmp/ReLeaSE/OpenChem')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openchem.models.Smiles2Label import Smiles2Label\n",
    "from openchem.modules.encoders.rnn_encoder import RNNEncoder\n",
    "from openchem.modules.mlp.openchem_mlp import OpenChemMLP\n",
    "from openchem.data.smiles_data_layer import SmilesDataset\n",
    "from openchem.data.utils import save_smiles_property_file\n",
    "from openchem.data.utils import create_loader\n",
    "from openchem.models.openchem_model import build_training, fit, evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "from torch.optim import Adam\n",
    "from torch.optim.lr_scheduler import ExponentialLR\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from h2o4gpu.metrics import r2_score, mean_squared_error\n",
    "import pandas as pd\n",
    "import copy\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openchem.data.utils import read_smiles_property_file\n",
    "data = read_smiles_property_file('./data/logP_labels.csv', \n",
    "                                 cols_to_read=[1, 2], keep_header=False)\n",
    "smiles = data[0]\n",
    "labels = data[1].astype('float32')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([1.300e+01, 1.420e+02, 9.030e+02, 3.624e+03, 5.417e+03, 2.995e+03,\n",
       "        7.490e+02, 2.700e+02, 5.800e+01, 5.000e+00]),\n",
       " array([-5.4       , -3.7310002 , -2.062     , -0.39300007,  1.2759999 ,\n",
       "         2.945     ,  4.614     ,  6.283     ,  7.952     ,  9.621     ,\n",
       "        11.29      ], dtype=float32),\n",
       " <a list of 10 Patch objects>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYAAAAD8CAYAAAB+UHOxAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMi4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvhp/UCwAAEZRJREFUeJzt3X2spGV5x/HvT/ClVSOLHJDugotx04h/qGSDtDaNFbMsaFyaSIJp6kZJNqaYaNKmrjUpVjSBNlVjUm2obLoaK1KVslEsblFj+gcvi/KOdhdE2S5lVxdRQ7TFXv1j7kOHZc7OnOWcmbPc308ymee5nntmruc5L7/zvMycVBWSpP48a9YNSJJmwwCQpE4ZAJLUKQNAkjplAEhSpwwASeqUASBJnTIAJKlTBoAkderYWTdwOCeccEKtXbt21m1I0lHl1ltv/XFVzY0bt6IDYO3atezatWvWbUjSUSXJDycZ5yEgSeqUASBJnTIAJKlTBoAkdcoAkKROGQCS1CkDQJI6ZQBIUqcMAEnq1Ip+J7A0ztqtX53Zaz9w2Ztm9trSUnAPQJI6ZQBIUqcMAEnqlAEgSZ0yACSpUwaAJHXKAJCkThkAktQpA0CSOmUASFKnDABJ6tREAZDkgSR3Jrktya5WOz7JziS72/2qVk+STyTZk+SOJGcMPc/mNn53ks3Ls0qSpEksZg/gD6rq1VW1vs1vBW6oqnXADW0e4FxgXbttAT4Fg8AALgFeC5wJXDIfGpKk6Xs6h4A2Advb9Hbg/KH6Z2rgRuC4JCcD5wA7q+pgVT0C7AQ2Po3XlyQ9DZMGQAFfT3Jrki2tdlJVPQTQ7k9s9dXAg0OP3dtqC9UlSTMw6f8DeF1V7UtyIrAzyfcOMzYjanWY+pMfPAiYLQCnnnrqhO1JkhZroj2AqtrX7vcD1zA4hv9wO7RDu9/fhu8FThl6+Bpg32Hqh77WFVW1vqrWz83NLW5tJEkTGxsASZ6f5IXz08AG4C5gBzB/Jc9m4No2vQN4e7sa6Czg0XaI6HpgQ5JV7eTvhlaTJM3AJIeATgKuSTI//p+q6l+T3AJcneQi4EfABW38dcB5wB7gMeAdAFV1MMmlwC1t3Ieq6uCSrYkkaVHGBkBV3Q+8akT9J8DZI+oFXLzAc20Dti2+TUnSUvOdwJLUKQNAkjplAEhSpwwASeqUASBJnTIAJKlTBoAkdcoAkKROGQCS1CkDQJI6ZQBIUqcMAEnqlAEgSZ0yACSpUwaAJHXKAJCkThkAktQpA0CSOmUASFKnDABJ6pQBIEmdMgAkqVMGgCR1ygCQpE4ZAJLUKQNAkjplAEhSpwwASeqUASBJnZo4AJIck+S7Sb7S5k9LclOS3Um+kOQ5rf7cNr+nLV879Bzvb/XvJzlnqVdGkjS5xewBvAe4d2j+cuBjVbUOeAS4qNUvAh6pqpcDH2vjSHI6cCHwSmAj8Mkkxzy99iVJR2qiAEiyBngT8Ok2H+ANwBfbkO3A+W16U5unLT+7jd8EXFVVv6qqHwB7gDOXYiUkSYs36R7Ax4E/B/63zb8Y+GlVPd7m9wKr2/Rq4EGAtvzRNv6J+ojHSJKmbGwAJHkzsL+qbh0ujxhaY5Yd7jHDr7clya4kuw4cODCuPUnSEZpkD+B1wFuSPABcxeDQz8eB45Ic28asAfa16b3AKQBt+YuAg8P1EY95QlVdUVXrq2r93NzcoldIkjSZsQFQVe+vqjVVtZbBSdxvVNUfAd8E3tqGbQaubdM72jxt+Teqqlr9wnaV0GnAOuDmJVsTSdKiHDt+yILeB1yV5MPAd4ErW/1K4LNJ9jD4y/9CgKq6O8nVwD3A48DFVfXrp/H6kqSnYVEBUFXfAr7Vpu9nxFU8VfVL4IIFHv8R4COLbVKStPR8J7AkdcoAkKROGQCS1CkDQJI6ZQBIUqcMAEnqlAEgSZ0yACSpUwaAJHXq6XwUhPSEtVu/OusWJC2SewCS1CkDQJI6ZQBIUqcMAEnqlAEgSZ0yACSpUwaAJHXKAJCkThkAktQpA0CSOmUASFKnDABJ6pQBIEmdMgAkqVMGgCR1ygCQpE4ZAJLUKQNAkjplAEhSp8YGQJLnJbk5ye1J7k7yV61+WpKbkuxO8oUkz2n157b5PW352qHnen+rfz/JOcu1UpKk8SbZA/gV8IaqehXwamBjkrOAy4GPVdU64BHgojb+IuCRqno58LE2jiSnAxcCrwQ2Ap9McsxSrowkaXJjA6AGftFmn91uBbwB+GKrbwfOb9Ob2jxt+dlJ0upXVdWvquoHwB7gzCVZC0nSok10DiDJMUluA/YDO4H7gJ9W1eNtyF5gdZteDTwI0JY/Crx4uD7iMZKkKTt2kkFV9Wvg1UmOA64BXjFqWLvPAssWqj9Jki3AFoBTTz11kvakmVi79aszed0HLnvTTF5XzzyLugqoqn4KfAs4CzguyXyArAH2tem9wCkAbfmLgIPD9RGPGX6NK6pqfVWtn5ubW0x7kqRFmOQqoLn2lz9JfgN4I3Av8E3grW3YZuDaNr2jzdOWf6OqqtUvbFcJnQasA25eqhWRJC3OJIeATga2tyt2ngVcXVVfSXIPcFWSDwPfBa5s468EPptkD4O//C8EqKq7k1wN3AM8DlzcDi1JkmZgbABU1R3Aa0bU72fEVTxV9UvgggWe6yPARxbfpiRpqflOYEnqlAEgSZ0yACSpUwaAJHXKAJCkThkAktQpA0CSOmUASFKnDABJ6pQBIEmdMgAkqVMGgCR1ygCQpE4ZAJLUKQNAkjplAEhSpwwASeqUASBJnTIAJKlTBoAkdcoAkKROGQCS1CkDQJI6ZQBIUqcMAEnqlAEgSZ0yACSpUwaAJHXKAJCkThkAktSpsQGQ5JQk30xyb5K7k7yn1Y9PsjPJ7na/qtWT5BNJ9iS5I8kZQ8+1uY3fnWTz8q2WJGmcSfYAHgf+tKpeAZwFXJzkdGArcENVrQNuaPMA5wLr2m0L8CkYBAZwCfBa4EzgkvnQkCRN39gAqKqHquo7bfrnwL3AamATsL0N2w6c36Y3AZ+pgRuB45KcDJwD7Kyqg1X1CLAT2LikayNJmtiizgEkWQu8BrgJOKmqHoJBSAAntmGrgQeHHra31RaqH/oaW5LsSrLrwIEDi2lPkrQIEwdAkhcAXwLeW1U/O9zQEbU6TP3Jhaorqmp9Va2fm5ubtD1J0iJNFABJns3gl//nqurLrfxwO7RDu9/f6nuBU4YevgbYd5i6JGkGJrkKKMCVwL1V9dGhRTuA+St5NgPXDtXf3q4GOgt4tB0iuh7YkGRVO/m7odUkSTNw7ARjXgf8MXBnktta7S+Ay4Crk1wE/Ai4oC27DjgP2AM8BrwDoKoOJrkUuKWN+1BVHVyStZAkLdrYAKiqf2f08XuAs0eML+DiBZ5rG7BtMQ1KkpaH7wSWpE4ZAJLUKQNAkjplAEhSpwwASeqUASBJnTIAJKlTBoAkdcoAkKROGQCS1CkDQJI6ZQBIUqcMAEnqlAEgSZ0yACSpUwaAJHXKAJCkThkAktQpA0CSOmUASFKnDABJ6pQBIEmdMgAkqVMGgCR1ygCQpE4ZAJLUKQNAkjplAEhSpwwASerU2ABIsi3J/iR3DdWOT7Izye52v6rVk+QTSfYkuSPJGUOP2dzG706yeXlWR5I0qUn2AP4R2HhIbStwQ1WtA25o8wDnAuvabQvwKRgEBnAJ8FrgTOCS+dCQJM3G2ACoqm8DBw8pbwK2t+ntwPlD9c/UwI3AcUlOBs4BdlbVwap6BNjJU0NFkjRFR3oO4KSqegig3Z/Y6quBB4fG7W21heqSpBlZ6pPAGVGrw9Sf+gTJliS7kuw6cODAkjYnSfp/RxoAD7dDO7T7/a2+FzhlaNwaYN9h6k9RVVdU1fqqWj83N3eE7UmSxjnSANgBzF/Jsxm4dqj+9nY10FnAo+0Q0fXAhiSr2snfDa0mSZqRY8cNSPJ54PXACUn2Mria5zLg6iQXAT8CLmjDrwPOA/YAjwHvAKiqg0kuBW5p4z5UVYeeWNbTtHbrV2fdgqSjyNgAqKq3LbDo7BFjC7h4gefZBmxbVHeSpGXjO4ElqVMGgCR1ygCQpE4ZAJLUqbEngSWtLLO82uuBy940s9fW0nMPQJI6ZQBIUqcMAEnqlAEgSZ0yACSpUwaAJHXKAJCkThkAktQpA0CSOmUASFKnDABJ6pQBIEmdMgAkqVMGgCR1ygCQpE4ZAJLUKQNAkjrlfwSTNLFZ/Tcy/xPZ8nAPQJI6ZQBIUqcMAEnqlAEgSZ3yJPAymNWJMklaDPcAJKlTUw+AJBuTfD/JniRbp/36kqSBqQZAkmOAvwPOBU4H3pbk9Gn2IEkamPY5gDOBPVV1P0CSq4BNwD1T7kPSUcQ3oC2PaQfAauDBofm9wGuX68U8GStJC5t2AGRErZ40INkCbGmzv0jy/WXq5QTgx8v03MvlaOwZ7HuajsaeYYX2ncsPu3hF9ty8dJJB0w6AvcApQ/NrgH3DA6rqCuCK5W4kya6qWr/cr7OUjsaewb6n6WjsGY7Ovo/Gng817auAbgHWJTktyXOAC4EdU+5BksSU9wCq6vEk7wauB44BtlXV3dPsQZI0MPV3AlfVdcB1037dEZb9MNMyOBp7BvuepqOxZzg6+z4ae36SVNX4UZKkZxw/CkKSOtVNACT5YJL/THJbu523wLgV81EVSf4myfeS3JHkmiTHLTDugSR3tvXaNe0+h/o47LZL8twkX2jLb0qydvpdPqmfU5J8M8m9Se5O8p4RY16f5NGh75u/nEWvhxr3Nc/AJ9q2viPJGbPo85CefntoO96W5GdJ3nvImJlv7yTbkuxPctdQ7fgkO5PsbverFnjs5jZmd5LN0+v6CFVVFzfgg8CfjRlzDHAf8DLgOcDtwOkz7HkDcGybvhy4fIFxDwAnzHj7jt12wJ8Af9+mLwS+MOOeTwbOaNMvBP5jRM+vB74yyz6P5GsOnAd8jcF7b84Cbpp1zyO+X/4LeOlK297A7wNnAHcN1f4a2Nqmt476WQSOB+5v96va9KpZb+vD3brZA5jQEx9VUVX/Dcx/VMVMVNXXq+rxNnsjg/dNrFSTbLtNwPY2/UXg7CSj3hw4FVX1UFV9p03/HLiXwbvVnwk2AZ+pgRuB45KcPOumhpwN3FdVP5x1I4eqqm8DBw8pD3/vbgfOH/HQc4CdVXWwqh4BdgIbl63RJdBbALy77Q5vW2AXbtRHVayUXwjvZPAX3SgFfD3Jre2d1LMwybZ7YkwLtkeBF0+luzHa4ajXADeNWPw7SW5P8rUkr5xqYwsb9zVfyd/LMNgD/PwCy1bi9j6pqh6CwR8OwIkjxqz0bf4Uz6h/CJPk34CXjFj0AeBTwKUMfnAuBf6WwS/VJz3FiMcu62VSh+u5qq5tYz4APA58boGneV1V7UtyIrAzyffaXzHTNMm2m/r2nUSSFwBfAt5bVT87ZPF3GBym+EU7b/QvwLpp9zjCuK/5itzWAO1NoG8B3j9i8Urd3pNYsdt8Ic+oAKiqN04yLsk/AF8ZsWjsR1UstXE9txNJbwbOrnagccRz7Gv3+5Ncw+BwzLQDYJJtNz9mb5JjgRfx1F3tqUrybAa//D9XVV8+dPlwIFTVdUk+meSEqprpZ8BM8DWf+vfyIpwLfKeqHj50wUrd3sDDSU6uqofaobT9I8bsZXAOY94a4FtT6O2IdXMI6JDjn38I3DVi2Ir6qIokG4H3AW+pqscWGPP8JC+cn2Zw4njUui23SbbdDmD+yoi3At9YKNSmoZ1/uBK4t6o+usCYl8yfp0hyJoOfmZ9Mr8uRPU3yNd8BvL1dDXQW8Oj8IYwV4G0scPhnJW7vZvh7dzNw7Ygx1wMbkqxqh5g3tNrKNeuz0NO6AZ8F7gTuYPDFPLnVfwu4bmjceQyuBrmPwWGYWfa8h8Exxdvabf4Kmid6ZnDVze3tdvcsex617YAPMQgwgOcB/9zW62bgZTPevr/HYBf9jqFtfB7wLuBdbcy723a9ncGJ+N+dZc+H+5of0ncY/POl+9r3/fpZ9936+k0Gv9BfNFRbUdubQTg9BPwPg7/qL2JwruoGYHe7P76NXQ98euix72zf33uAd8x6e4+7+U5gSepUN4eAJElPZgBIUqcMAEnqlAEgSZ0yACSpUwaAJHXKAJCkThkAktSp/wM39h/2JvgJ3wAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "plt.hist(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openchem.data.utils import get_tokens\n",
    "tokens, _, _ = get_tokens(smiles)\n",
    "tokens = ['<', '>', '#', '%', ')', '(', '+', '-', '/', '.', '1', '0', '3', '2', '5', '4', '7',\n",
    "          '6', '9', '8', '=', 'A', '@', 'C', 'B', 'F', 'I', 'H', 'O', 'N', 'P', 'S', '[', ']',\n",
    "          '\\\\', 'c', 'e', 'i', 'l', 'o', 'n', 'p', 's', 'r', '\\n']\n",
    "tokens = ''.join(tokens) + ' '"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we define the architecture of our Recurrent Neural Network (RNN). This model solve the regression problem, i.e. predicts values of logP directly from SMILES (no need to calculated descriptors). \n",
    "\n",
    "We will use 2 unidirectional LSTM layers with hidden size of 128 each. Embedding size is 128. Dense layers is a Multi-layer Perceptron with with hidden size of 128 and ReLU activation function. The loss function is simple MSE (mean squared error).\n",
    "\n",
    "We will use OpenChem toolkit to train the model. For more details on OpenChem visit: https://mariewelt.github.io/OpenChem/\n",
    "\n",
    "The architecture of the network is shown below:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"./figures/predictor_figure.png\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from openchem.utils.utils import identity\n",
    "from openchem.modules.embeddings.basic_embedding import Embedding\n",
    "model_object = Smiles2Label\n",
    "\n",
    "n_hidden = 128\n",
    "batch_size = 128\n",
    "num_epochs = 26\n",
    "lr = 0.005\n",
    "\n",
    "model_params = {\n",
    "    'use_cuda': True,\n",
    "    'random_seed': 42,\n",
    "    'world_size': 1,\n",
    "    'task': 'regression',\n",
    "    'data_layer': SmilesDataset,\n",
    "    'use_clip_grad': False,\n",
    "    'batch_size': batch_size,\n",
    "    'num_epochs': num_epochs,\n",
    "    'logdir': './checkpoints/logP/',\n",
    "    'print_every': 1,\n",
    "    'save_every': 5,\n",
    "    'train_data_layer': None,\n",
    "    'val_data_layer': None,\n",
    "    'eval_metrics': r2_score,\n",
    "    'criterion': nn.MSELoss(),\n",
    "    'optimizer': Adam,\n",
    "    'optimizer_params': {\n",
    "        'lr': lr,\n",
    "    },\n",
    "    'lr_scheduler': ExponentialLR,\n",
    "    'lr_scheduler_params': {\n",
    "        'gamma': 0.98\n",
    "    },\n",
    "    'embedding': Embedding,\n",
    "    'embedding_params': {\n",
    "        'num_embeddings': len(tokens),\n",
    "        'embedding_dim': n_hidden,\n",
    "        'padding_idx': tokens.index(' ')\n",
    "    },\n",
    "    'encoder': RNNEncoder,\n",
    "    'encoder_params': {\n",
    "        'input_size': n_hidden,\n",
    "        'layer': \"LSTM\",\n",
    "        'encoder_dim': n_hidden,\n",
    "        'n_layers': 2,\n",
    "        'dropout': 0.8,\n",
    "        'is_bidirectional': False\n",
    "    },\n",
    "    'mlp': OpenChemMLP,\n",
    "    'mlp_params': {\n",
    "        'input_size': n_hidden,\n",
    "        'n_layers': 2,\n",
    "        'hidden_size': [n_hidden, 1],\n",
    "        'activation': [F.relu, identity],\n",
    "        'dropout': 0.0\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "pickle.dump(model_params, open('./checkpoints/logP/model_parameters.pkl', 'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    os.stat(model_params['logdir'])\n",
    "except:\n",
    "    os.mkdir(model_params['logdir'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_dir = './checkpoints/logP/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_data_dir = './data/tmp/'\n",
    "try:\n",
    "    os.stat(tmp_data_dir)\n",
    "except:\n",
    "    os.mkdir(tmp_data_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initializing data splitter for cross validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from h2o4gpu.model_selection import KFold, train_test_split\n",
    "\n",
    "cross_validation_split = KFold(n_splits=5, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = list(cross_validation_split.split(smiles, labels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training cross-validated models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cross validation, fold number 0 in progress...\n",
      "TRAINING: [Time: 0m 3s, Epoch: 0, Progress: 0%, Loss: 2.3622]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 1.2985, Metrics: 0.6359]\n",
      "TRAINING: [Time: 0m 7s, Epoch: 1, Progress: 3%, Loss: 1.0032]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.8611, Metrics: 0.7517]\n",
      "TRAINING: [Time: 0m 11s, Epoch: 2, Progress: 7%, Loss: 0.7285]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.5726, Metrics: 0.8392]\n",
      "TRAINING: [Time: 0m 15s, Epoch: 3, Progress: 11%, Loss: 0.5623]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.6304, Metrics: 0.8219]\n",
      "TRAINING: [Time: 0m 20s, Epoch: 4, Progress: 15%, Loss: 0.5127]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.5067, Metrics: 0.8545]\n",
      "TRAINING: [Time: 0m 24s, Epoch: 5, Progress: 19%, Loss: 0.4698]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4158, Metrics: 0.8812]\n",
      "TRAINING: [Time: 0m 28s, Epoch: 6, Progress: 23%, Loss: 0.4491]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4470, Metrics: 0.8731]\n",
      "TRAINING: [Time: 0m 32s, Epoch: 7, Progress: 26%, Loss: 0.4129]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4512, Metrics: 0.8701]\n",
      "TRAINING: [Time: 0m 36s, Epoch: 8, Progress: 30%, Loss: 0.3738]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4219, Metrics: 0.8775]\n",
      "TRAINING: [Time: 0m 40s, Epoch: 9, Progress: 34%, Loss: 0.3544]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3709, Metrics: 0.8945]\n",
      "TRAINING: [Time: 0m 44s, Epoch: 10, Progress: 38%, Loss: 0.3496]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4238, Metrics: 0.8786]\n",
      "TRAINING: [Time: 0m 48s, Epoch: 11, Progress: 42%, Loss: 0.3387]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3569, Metrics: 0.8985]\n",
      "TRAINING: [Time: 0m 52s, Epoch: 12, Progress: 46%, Loss: 0.3206]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3086, Metrics: 0.9114]\n",
      "TRAINING: [Time: 0m 56s, Epoch: 13, Progress: 50%, Loss: 0.2972]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3252, Metrics: 0.9061]\n",
      "TRAINING: [Time: 1m 0s, Epoch: 14, Progress: 53%, Loss: 0.2924]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3580, Metrics: 0.8974]\n",
      "TRAINING: [Time: 1m 4s, Epoch: 15, Progress: 57%, Loss: 0.2834]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3341, Metrics: 0.9037]\n",
      "TRAINING: [Time: 1m 8s, Epoch: 16, Progress: 61%, Loss: 0.2688]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3273, Metrics: 0.9060]\n",
      "TRAINING: [Time: 1m 13s, Epoch: 17, Progress: 65%, Loss: 0.2653]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2951, Metrics: 0.9151]\n",
      "TRAINING: [Time: 1m 17s, Epoch: 18, Progress: 69%, Loss: 0.2616]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2984, Metrics: 0.9135]\n",
      "TRAINING: [Time: 1m 21s, Epoch: 19, Progress: 73%, Loss: 0.2533]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2794, Metrics: 0.9195]\n",
      "TRAINING: [Time: 1m 26s, Epoch: 20, Progress: 76%, Loss: 0.2310]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2958, Metrics: 0.9155]\n",
      "TRAINING: [Time: 1m 30s, Epoch: 21, Progress: 80%, Loss: 0.2278]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2782, Metrics: 0.9204]\n",
      "TRAINING: [Time: 1m 35s, Epoch: 22, Progress: 84%, Loss: 0.2207]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2739, Metrics: 0.9215]\n",
      "TRAINING: [Time: 1m 39s, Epoch: 23, Progress: 88%, Loss: 0.2236]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2859, Metrics: 0.9182]\n",
      "TRAINING: [Time: 1m 44s, Epoch: 24, Progress: 92%, Loss: 0.2213]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3009, Metrics: 0.9140]\n",
      "TRAINING: [Time: 1m 48s, Epoch: 25, Progress: 96%, Loss: 0.2040]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3353, Metrics: 0.9043]\n",
      "Cross validation, fold number 1 in progress...\n",
      "TRAINING: [Time: 0m 3s, Epoch: 0, Progress: 0%, Loss: 2.1325]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 1.1821, Metrics: 0.6189]\n",
      "TRAINING: [Time: 0m 8s, Epoch: 1, Progress: 3%, Loss: 0.9203]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.8710, Metrics: 0.7209]\n",
      "TRAINING: [Time: 0m 12s, Epoch: 2, Progress: 7%, Loss: 0.7350]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4785, Metrics: 0.8480]\n",
      "TRAINING: [Time: 0m 16s, Epoch: 3, Progress: 11%, Loss: 0.5794]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.5023, Metrics: 0.8385]\n",
      "TRAINING: [Time: 0m 21s, Epoch: 4, Progress: 15%, Loss: 0.5058]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.6194, Metrics: 0.8015]\n",
      "TRAINING: [Time: 0m 25s, Epoch: 5, Progress: 19%, Loss: 0.4898]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4432, Metrics: 0.8570]\n",
      "TRAINING: [Time: 0m 30s, Epoch: 6, Progress: 23%, Loss: 0.4307]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.5212, Metrics: 0.8332]\n",
      "TRAINING: [Time: 0m 34s, Epoch: 7, Progress: 26%, Loss: 0.4014]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3416, Metrics: 0.8895]\n",
      "TRAINING: [Time: 0m 39s, Epoch: 8, Progress: 30%, Loss: 0.3837]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4171, Metrics: 0.8661]\n",
      "TRAINING: [Time: 0m 43s, Epoch: 9, Progress: 34%, Loss: 0.3694]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3545, Metrics: 0.8860]\n",
      "TRAINING: [Time: 0m 47s, Epoch: 10, Progress: 38%, Loss: 0.3582]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3321, Metrics: 0.8940]\n",
      "TRAINING: [Time: 0m 52s, Epoch: 11, Progress: 42%, Loss: 0.3344]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3363, Metrics: 0.8912]\n",
      "TRAINING: [Time: 0m 56s, Epoch: 12, Progress: 46%, Loss: 0.3205]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4540, Metrics: 0.8546]\n",
      "TRAINING: [Time: 1m 1s, Epoch: 13, Progress: 50%, Loss: 0.2994]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4455, Metrics: 0.8589]\n",
      "TRAINING: [Time: 1m 5s, Epoch: 14, Progress: 53%, Loss: 0.2885]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3040, Metrics: 0.9027]\n",
      "TRAINING: [Time: 1m 9s, Epoch: 15, Progress: 57%, Loss: 0.2806]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4661, Metrics: 0.8511]\n",
      "TRAINING: [Time: 1m 14s, Epoch: 16, Progress: 61%, Loss: 0.2634]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2926, Metrics: 0.9062]\n",
      "TRAINING: [Time: 1m 18s, Epoch: 17, Progress: 65%, Loss: 0.2632]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3383, Metrics: 0.8918]\n",
      "TRAINING: [Time: 1m 23s, Epoch: 18, Progress: 69%, Loss: 0.2604]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2939, Metrics: 0.9040]\n",
      "TRAINING: [Time: 1m 27s, Epoch: 19, Progress: 73%, Loss: 0.2400]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3151, Metrics: 0.8996]\n",
      "TRAINING: [Time: 1m 32s, Epoch: 20, Progress: 76%, Loss: 0.2388]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2703, Metrics: 0.9132]\n",
      "TRAINING: [Time: 1m 36s, Epoch: 21, Progress: 80%, Loss: 0.2395]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2847, Metrics: 0.9091]\n",
      "TRAINING: [Time: 1m 40s, Epoch: 22, Progress: 84%, Loss: 0.2250]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2599, Metrics: 0.9167]\n",
      "TRAINING: [Time: 1m 45s, Epoch: 23, Progress: 88%, Loss: 0.2151]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2939, Metrics: 0.9055]\n",
      "TRAINING: [Time: 1m 49s, Epoch: 24, Progress: 92%, Loss: 0.2116]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2668, Metrics: 0.9138]\n",
      "TRAINING: [Time: 1m 54s, Epoch: 25, Progress: 96%, Loss: 0.2065]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2759, Metrics: 0.9111]\n",
      "Cross validation, fold number 2 in progress...\n",
      "TRAINING: [Time: 0m 3s, Epoch: 0, Progress: 0%, Loss: 1.8440]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 1.3091, Metrics: 0.6024]\n",
      "TRAINING: [Time: 0m 8s, Epoch: 1, Progress: 3%, Loss: 0.8846]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.8430, Metrics: 0.7473]\n",
      "TRAINING: [Time: 0m 12s, Epoch: 2, Progress: 7%, Loss: 0.6331]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.5139, Metrics: 0.8443]\n",
      "TRAINING: [Time: 0m 16s, Epoch: 3, Progress: 11%, Loss: 0.5476]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.5639, Metrics: 0.8289]\n",
      "TRAINING: [Time: 0m 21s, Epoch: 4, Progress: 15%, Loss: 0.4902]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4199, Metrics: 0.8742]\n",
      "TRAINING: [Time: 0m 25s, Epoch: 5, Progress: 19%, Loss: 0.4437]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.6000, Metrics: 0.8180]\n",
      "TRAINING: [Time: 0m 30s, Epoch: 6, Progress: 23%, Loss: 0.4192]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4403, Metrics: 0.8667]\n",
      "TRAINING: [Time: 0m 34s, Epoch: 7, Progress: 26%, Loss: 0.3785]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3311, Metrics: 0.9010]\n",
      "TRAINING: [Time: 0m 38s, Epoch: 8, Progress: 30%, Loss: 0.3793]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.5808, Metrics: 0.8246]\n",
      "TRAINING: [Time: 0m 43s, Epoch: 9, Progress: 34%, Loss: 0.3582]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4701, Metrics: 0.8593]\n",
      "TRAINING: [Time: 0m 47s, Epoch: 10, Progress: 38%, Loss: 0.3313]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3272, Metrics: 0.9016]\n",
      "TRAINING: [Time: 0m 51s, Epoch: 11, Progress: 42%, Loss: 0.3160]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3314, Metrics: 0.9001]\n",
      "TRAINING: [Time: 0m 56s, Epoch: 12, Progress: 46%, Loss: 0.2908]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3095, Metrics: 0.9078]\n",
      "TRAINING: [Time: 1m 0s, Epoch: 13, Progress: 50%, Loss: 0.2800]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2800, Metrics: 0.9161]\n",
      "TRAINING: [Time: 1m 5s, Epoch: 14, Progress: 53%, Loss: 0.2715]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EVALUATION: [Time: 0m 0s, Loss: 0.3136, Metrics: 0.9062]\n",
      "TRAINING: [Time: 1m 9s, Epoch: 15, Progress: 57%, Loss: 0.2682]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3307, Metrics: 0.9004]\n",
      "TRAINING: [Time: 1m 13s, Epoch: 16, Progress: 61%, Loss: 0.2546]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3475, Metrics: 0.8947]\n",
      "TRAINING: [Time: 1m 18s, Epoch: 17, Progress: 65%, Loss: 0.2429]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2908, Metrics: 0.9126]\n",
      "TRAINING: [Time: 1m 22s, Epoch: 18, Progress: 69%, Loss: 0.2320]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3450, Metrics: 0.8972]\n",
      "TRAINING: [Time: 1m 27s, Epoch: 19, Progress: 73%, Loss: 0.2424]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3765, Metrics: 0.8868]\n",
      "TRAINING: [Time: 1m 31s, Epoch: 20, Progress: 76%, Loss: 0.2462]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4004, Metrics: 0.8789]\n",
      "TRAINING: [Time: 1m 35s, Epoch: 21, Progress: 80%, Loss: 0.2327]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3599, Metrics: 0.8912]\n",
      "TRAINING: [Time: 1m 40s, Epoch: 22, Progress: 84%, Loss: 0.2177]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2877, Metrics: 0.9137]\n",
      "TRAINING: [Time: 1m 44s, Epoch: 23, Progress: 88%, Loss: 0.2110]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3199, Metrics: 0.9035]\n",
      "TRAINING: [Time: 1m 48s, Epoch: 24, Progress: 92%, Loss: 0.2002]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3813, Metrics: 0.8850]\n",
      "TRAINING: [Time: 1m 53s, Epoch: 25, Progress: 96%, Loss: 0.2022]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2826, Metrics: 0.9148]\n",
      "Cross validation, fold number 3 in progress...\n",
      "TRAINING: [Time: 0m 3s, Epoch: 0, Progress: 0%, Loss: 1.9366]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 1.4129, Metrics: 0.6009]\n",
      "TRAINING: [Time: 0m 8s, Epoch: 1, Progress: 3%, Loss: 0.8814]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.7201, Metrics: 0.7947]\n",
      "TRAINING: [Time: 0m 12s, Epoch: 2, Progress: 7%, Loss: 0.6631]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 1.0425, Metrics: 0.6997]\n",
      "TRAINING: [Time: 0m 16s, Epoch: 3, Progress: 11%, Loss: 0.5728]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.9165, Metrics: 0.7397]\n",
      "TRAINING: [Time: 0m 20s, Epoch: 4, Progress: 15%, Loss: 0.4986]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.5831, Metrics: 0.8359]\n",
      "TRAINING: [Time: 0m 25s, Epoch: 5, Progress: 19%, Loss: 0.4665]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 1.0738, Metrics: 0.6948]\n",
      "TRAINING: [Time: 0m 29s, Epoch: 6, Progress: 23%, Loss: 0.4390]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4211, Metrics: 0.8788]\n",
      "TRAINING: [Time: 0m 33s, Epoch: 7, Progress: 26%, Loss: 0.4167]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4656, Metrics: 0.8678]\n",
      "TRAINING: [Time: 0m 38s, Epoch: 8, Progress: 30%, Loss: 0.3866]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4054, Metrics: 0.8825]\n",
      "TRAINING: [Time: 0m 42s, Epoch: 9, Progress: 34%, Loss: 0.3582]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4479, Metrics: 0.8702]\n",
      "TRAINING: [Time: 0m 46s, Epoch: 10, Progress: 38%, Loss: 0.3441]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4158, Metrics: 0.8815]\n",
      "TRAINING: [Time: 0m 51s, Epoch: 11, Progress: 42%, Loss: 0.3240]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3989, Metrics: 0.8855]\n",
      "TRAINING: [Time: 0m 55s, Epoch: 12, Progress: 46%, Loss: 0.3186]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.6286, Metrics: 0.8182]\n",
      "TRAINING: [Time: 0m 59s, Epoch: 13, Progress: 50%, Loss: 0.2970]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3353, Metrics: 0.9050]\n",
      "TRAINING: [Time: 1m 4s, Epoch: 14, Progress: 53%, Loss: 0.2868]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4000, Metrics: 0.8833]\n",
      "TRAINING: [Time: 1m 8s, Epoch: 15, Progress: 57%, Loss: 0.2914]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4039, Metrics: 0.8853]\n",
      "TRAINING: [Time: 1m 12s, Epoch: 16, Progress: 61%, Loss: 0.2717]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3432, Metrics: 0.9028]\n",
      "TRAINING: [Time: 1m 16s, Epoch: 17, Progress: 65%, Loss: 0.2619]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4046, Metrics: 0.8858]\n",
      "TRAINING: [Time: 1m 21s, Epoch: 18, Progress: 69%, Loss: 0.2506]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3248, Metrics: 0.9074]\n",
      "TRAINING: [Time: 1m 25s, Epoch: 19, Progress: 73%, Loss: 0.2428]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4984, Metrics: 0.8581]\n",
      "TRAINING: [Time: 1m 29s, Epoch: 20, Progress: 76%, Loss: 0.2450]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3104, Metrics: 0.9122]\n",
      "TRAINING: [Time: 1m 34s, Epoch: 21, Progress: 80%, Loss: 0.2319]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4155, Metrics: 0.8832]\n",
      "TRAINING: [Time: 1m 38s, Epoch: 22, Progress: 84%, Loss: 0.2228]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2810, Metrics: 0.9197]\n",
      "TRAINING: [Time: 1m 42s, Epoch: 23, Progress: 88%, Loss: 0.2223]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3978, Metrics: 0.8863]\n",
      "TRAINING: [Time: 1m 47s, Epoch: 24, Progress: 92%, Loss: 0.2249]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2956, Metrics: 0.9157]\n",
      "TRAINING: [Time: 1m 51s, Epoch: 25, Progress: 96%, Loss: 0.2143]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4066, Metrics: 0.8843]\n",
      "Cross validation, fold number 4 in progress...\n",
      "TRAINING: [Time: 0m 3s, Epoch: 0, Progress: 0%, Loss: 1.9487]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 1.1392, Metrics: 0.6495]\n",
      "TRAINING: [Time: 0m 6s, Epoch: 1, Progress: 3%, Loss: 0.9011]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.6441, Metrics: 0.8052]\n",
      "TRAINING: [Time: 0m 10s, Epoch: 2, Progress: 7%, Loss: 0.6240]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4907, Metrics: 0.8503]\n",
      "TRAINING: [Time: 0m 13s, Epoch: 3, Progress: 11%, Loss: 0.5284]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4707, Metrics: 0.8557]\n",
      "TRAINING: [Time: 0m 17s, Epoch: 4, Progress: 15%, Loss: 0.5058]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4693, Metrics: 0.8543]\n",
      "TRAINING: [Time: 0m 20s, Epoch: 5, Progress: 19%, Loss: 0.4418]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.6393, Metrics: 0.8014]\n",
      "TRAINING: [Time: 0m 24s, Epoch: 6, Progress: 23%, Loss: 0.4222]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3912, Metrics: 0.8789]\n",
      "TRAINING: [Time: 0m 27s, Epoch: 7, Progress: 26%, Loss: 0.3896]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4188, Metrics: 0.8708]\n",
      "TRAINING: [Time: 0m 31s, Epoch: 8, Progress: 30%, Loss: 0.3649]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4376, Metrics: 0.8646]\n",
      "TRAINING: [Time: 0m 34s, Epoch: 9, Progress: 34%, Loss: 0.3449]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3713, Metrics: 0.8853]\n",
      "TRAINING: [Time: 0m 38s, Epoch: 10, Progress: 38%, Loss: 0.3265]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3598, Metrics: 0.8896]\n",
      "TRAINING: [Time: 0m 41s, Epoch: 11, Progress: 42%, Loss: 0.3136]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3528, Metrics: 0.8901]\n",
      "TRAINING: [Time: 0m 45s, Epoch: 12, Progress: 46%, Loss: 0.3050]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3395, Metrics: 0.8957]\n",
      "TRAINING: [Time: 0m 48s, Epoch: 13, Progress: 50%, Loss: 0.3029]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2968, Metrics: 0.9075]\n",
      "TRAINING: [Time: 0m 52s, Epoch: 14, Progress: 53%, Loss: 0.2781]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4025, Metrics: 0.8751]\n",
      "TRAINING: [Time: 0m 56s, Epoch: 15, Progress: 57%, Loss: 0.2716]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3135, Metrics: 0.9035]\n",
      "TRAINING: [Time: 0m 59s, Epoch: 16, Progress: 61%, Loss: 0.2509]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3666, Metrics: 0.8875]\n",
      "TRAINING: [Time: 1m 3s, Epoch: 17, Progress: 65%, Loss: 0.2453]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3025, Metrics: 0.9064]\n",
      "TRAINING: [Time: 1m 6s, Epoch: 18, Progress: 69%, Loss: 0.2461]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3154, Metrics: 0.9033]\n",
      "TRAINING: [Time: 1m 10s, Epoch: 19, Progress: 73%, Loss: 0.2416]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.5205, Metrics: 0.8389]\n",
      "TRAINING: [Time: 1m 13s, Epoch: 20, Progress: 76%, Loss: 0.2256]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2865, Metrics: 0.9115]\n",
      "TRAINING: [Time: 1m 17s, Epoch: 21, Progress: 80%, Loss: 0.2309]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2951, Metrics: 0.9097]\n",
      "TRAINING: [Time: 1m 20s, Epoch: 22, Progress: 84%, Loss: 0.2157]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4373, Metrics: 0.8645]\n",
      "TRAINING: [Time: 1m 24s, Epoch: 23, Progress: 88%, Loss: 0.2069]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3164, Metrics: 0.9019]\n",
      "TRAINING: [Time: 1m 27s, Epoch: 24, Progress: 92%, Loss: 0.2047]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2896, Metrics: 0.9103]\n",
      "TRAINING: [Time: 1m 31s, Epoch: 25, Progress: 96%, Loss: 0.2085]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3524, Metrics: 0.8927]\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "i = 0\n",
    "models = []\n",
    "results = []\n",
    "for split in data:\n",
    "    print('Cross validation, fold number ' + str(i) + ' in progress...')\n",
    "    train, test = split\n",
    "    X_train = smiles[train]\n",
    "    y_train = labels[train].reshape(-1)\n",
    "    X_test = smiles[test]\n",
    "    y_test = labels[test].reshape(-1)\n",
    "    save_smiles_property_file(tmp_data_dir + str(i) + '_train.smi', \n",
    "                              X_train, y_train.reshape(-1, 1))\n",
    "    save_smiles_property_file(tmp_data_dir + str(i) + '_test.smi', \n",
    "                              X_test, y_test.reshape(-1, 1))\n",
    "\n",
    "    train_dataset = SmilesDataset(tmp_data_dir + str(i) + '_train.smi',\n",
    "                           delimiter=',', cols_to_read=[0, 1], tokens=tokens,\n",
    "                                 flip=False,)\n",
    "    train_dataset.target = train_dataset.target\n",
    "    test_dataset = SmilesDataset(tmp_data_dir + str(i) + '_test.smi',\n",
    "                       delimiter=',', cols_to_read=[0, 1], tokens=tokens,\n",
    "                                flip=False)\n",
    "    test_dataset.target = test_dataset.target\n",
    "    model_params['train_data_layer'] = train_dataset\n",
    "    model_params['val_data_layer'] = test_dataset\n",
    "    model_params['logdir'] = log_dir + 'fold_' + str(i)  \n",
    "    ckpt_dir = model_params['logdir'] + '/checkpoint/'\n",
    "    try:\n",
    "        os.stat(ckpt_dir)\n",
    "    except:\n",
    "        os.mkdir(model_params['logdir'])\n",
    "        os.mkdir(ckpt_dir)\n",
    "    train_loader = create_loader(train_dataset,\n",
    "                             batch_size=model_params['batch_size'],\n",
    "                             shuffle=True,\n",
    "                             num_workers=4,\n",
    "                             pin_memory=True,\n",
    "                             sampler=None)\n",
    "    val_loader = create_loader(test_dataset,\n",
    "                           batch_size=model_params['batch_size'],\n",
    "                           shuffle=False,\n",
    "                           num_workers=1,\n",
    "                           pin_memory=True)\n",
    "    models.append(model_object(params=model_params).cuda())\n",
    "    criterion, optimizer, lr_scheduler = build_training(models[i], model_params)\n",
    "    results.append(fit(models[i], lr_scheduler, train_loader, optimizer, criterion,\n",
    "        model_params, eval=True, val_loader=val_loader))\n",
    "    \n",
    "    i = i+1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating the models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EVALUATION: [Time: 0m 0s, Loss: 0.3353, Metrics: 0.9043]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2759, Metrics: 0.9111]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.2826, Metrics: 0.9148]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.4066, Metrics: 0.8843]\n",
      "EVALUATION: [Time: 0m 0s, Loss: 0.3524, Metrics: 0.8927]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "rmse = []\n",
    "auc_score = []\n",
    "for i in range(5):\n",
    "    test_dataset = SmilesDataset(tmp_data_dir + str(i) + '_test.smi',\n",
    "                                 delimiter=',', cols_to_read=[0, 1], tokens=tokens,\n",
    "                                flip=False)\n",
    "    test_dataset.target = test_dataset.target\n",
    "    val_loader = create_loader(test_dataset,\n",
    "                               batch_size=model_params['batch_size'],\n",
    "                               shuffle=False,\n",
    "                               num_workers=1,\n",
    "                               pin_memory=True)\n",
    "    metrics = evaluate(models[i], val_loader, criterion)\n",
    "    rmse.append(np.sqrt(metrics[0]))\n",
    "    auc_score.append(metrics[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cross-validated RMSE:  0.5734339993856861\n",
      "Cross-validated R2 score:  0.9014448637134084\n"
     ]
    }
   ],
   "source": [
    "print(\"Cross-validated RMSE: \",  np.mean(rmse))\n",
    "print(\"Cross-validated R2 score: \", np.mean(auc_score))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mv ./checkpoints/logP/fold_0/checkpoint/epoch_25 ./checkpoints/logP/fold_0.pkl\n",
    "! mv ./checkpoints/logP/fold_1/checkpoint/epoch_25 ./checkpoints/logP/fold_1.pkl\n",
    "! mv ./checkpoints/logP/fold_2/checkpoint/epoch_25 ./checkpoints/logP/fold_2.pkl\n",
    "! mv ./checkpoints/logP/fold_3/checkpoint/epoch_25 ./checkpoints/logP/fold_3.pkl\n",
    "! mv ./checkpoints/logP/fold_4/checkpoint/epoch_25 ./checkpoints/logP/fold_4.pkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
